
from PyQt4.QtGui import QIcon, QMainWindow, QCursor
from PyQt4.QtCore import QObject, Qt, SIGNAL

try:
    from matplotlib.backends import backend_qt4
    import matplotlib
    
    class SCFigureManagerQT(backend_qt4.FigureManagerQT):
        """
        Patching matplotlib...
        """
        def __init__(self, canvas, num):
            if backend_qt4.DEBUG:
                print 'FigureManagerQT.%s' % backend_qt4.fn_name()
            backend_qt4.FigureManagerBase.__init__(self, canvas, num)
            self.canvas = canvas
            
            self.window = QMainWindow()
            self.window.setWindowTitle("Figure %d" % num)
            self.window.setAttribute(Qt.WA_DeleteOnClose)
    
    #        image = os.path.join(matplotlib.rcParams['datapath'],
    #                         'images', 'matplotlib.png' )
            self.window.setWindowIcon(QIcon(":/Icons/BaseIcon2_128.png"))
    
            # Give the keyboard focus to the figure instead of the manager
            self.canvas.setFocusPolicy(Qt.ClickFocus)
            self.canvas.setFocus()
    
            QObject.connect(self.window, SIGNAL('destroyed()'),
                            self._widgetclosed)
            self.window._destroying = False
    
            self.toolbar = self._get_toolbar(self.canvas, self.window)
            self.window.addToolBar(self.toolbar)
            QObject.connect(self.toolbar, SIGNAL("message"),
                    self.window.statusBar().showMessage)
    
            self.window.setCentralWidget(self.canvas)
    
            if matplotlib.is_interactive():
                self.window.show()
    
            # attach a show method to the figure for pylab ease of use
            self.canvas.figure.show = lambda *args: self.window.show()
            
            self.canvas.axes = self.canvas.figure.add_subplot(111)
    
            def notify_axes_change( fig ):
                # This will be called whenever the current axes is changed
                if self.toolbar != None: self.toolbar.update()
            self.canvas.figure.add_axobserver(notify_axes_change)
    
    class SCNavigationToolbar2QT( backend_qt4.NavigationToolbar2QT ):
        def _init_toolbar(self):
            super(SCNavigationToolbar2QT, self)._init_toolbar()
            
        def set_cursor( self, cursor ):
            if backend_qt4.DEBUG: print 'Set cursor' , cursor
            self.parent().setCursor( QCursor(backend_qt4.cursord[cursor]) )
            
        def _icon(self, name):
            if 'back' in name:
                return QIcon(":/Icons/Actions/Plot/back.svg")
            if 'forward' in name:
                return QIcon(":/Icons/Actions/Plot/forward.svg")
            if 'move' in name:
                return QIcon(":/Icons/Actions/Plot/move.svg")
            if 'home' in name:
                return QIcon(":/Icons/Actions/Plot/home.svg")
            if 'filesave' in name:
                return QIcon(":/Icons/Actions/Plot/filesave.svg")
            if 'subplot' in name:
                return QIcon(":/Icons/Actions/Plot/subplots.svg")
            if 'zoom_to_rect' in name:
                return QIcon(":/Icons/Actions/Plot/zoom_to_rect.svg")

    from mpl_toolkits.mplot3d import Axes3D
    import matplotlib.pyplot as plt
    
#    def plot(x,y=None,fignum=1,labels=('',''),legend=False,legpos='best',*args,**kwargs):
#        '''
#        '''
#        plt.figure(fignum)
#        plt.clf()
#        if y:
#            plt.plot(x,y,*args,**kwargs)
#        else:
#            plt.plot(x,*args,**kwargs)
#        plt.xlabel(labels[0])
#        plt.ylabel(labels[1])
#        if legend: plt.legend(pos=legpos)
#        plt.show()
        
    def plot3d(x,y,z,ptype='scatter',fignum=1,labels=('','',''),**kwargs):
        '''
        ptype, plot type can be "scatter", "wireframe", "surface", "contour", "contourf"
        '''
        fig = plt.figure(fignum)
        plt.clf()
        ax = Axes3D(fig)
        if ptype == 'scatter':
            ax.scatter(x,y,z,**kwargs)
        elif ptype == 'wireframe':
            ax.plot_wireframe(x,y,z,**kwargs)
        elif ptype == 'surface':
            ax.plot_surface(x,y,z,**kwargs)
        elif ptype == 'contour':
            ax.contour(x,y,z,**kwargs)
        elif ptype == 'contourf':
            ax.contourf(x,y,z,**kwargs)
        else:
            raise NotImplementedError()
        ax.set_xlabel(labels[0])
        ax.set_ylabel(labels[1])
        ax.set_zlabel(labels[2])
        plt.show()
        return ax

except ImportError:
    pass